# -*- coding: utf-8 -*-
import os
from openpyxl import Workbook

from alibabacloud_slb20140515.client import Client as SlbClient
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_slb20140515 import models as slb_models


class SLBExporter:
    def __init__(self, region='ap-southeast-1'):
        self.region_id = region
        self.client = self.create_client()

    @staticmethod
    def create_client() -> SlbClient:
        config = open_api_models.Config(
            access_key_id=os.environ['ALIBABA_CLOUD_ACCESS_KEY_ID'],
            access_key_secret=os.environ['ALIBABA_CLOUD_ACCESS_KEY_SECRET']
        )
        config.endpoint = 'slb.ap-southeast-1.aliyuncs.com'
        return SlbClient(config)

    def list_load_balancers(self):
        req = slb_models.DescribeLoadBalancersRequest(region_id=self.region_id)
        resp = self.client.describe_load_balancers(req)
        return resp.body.load_balancers.load_balancer

    def get_listener_ports_and_ip(self, lb_id):
        req = slb_models.DescribeLoadBalancerAttributeRequest(load_balancer_id=lb_id)
        resp = self.client.describe_load_balancer_attribute(req)
        ip = resp.body.address  # 🔥 加载均衡 IP
        listeners = resp.body.listener_ports_and_protocol.listener_port_and_protocol
        return listeners, ip

    def get_listener_attr(self, lb_id, port, protocol):
        protocol = protocol.lower()
        request_cls = {
            'http': slb_models.DescribeLoadBalancerHTTPListenerAttributeRequest,
            'https': slb_models.DescribeLoadBalancerHTTPSListenerAttributeRequest,
            'tcp': slb_models.DescribeLoadBalancerTCPListenerAttributeRequest,
            'udp': slb_models.DescribeLoadBalancerUDPListenerAttributeRequest,
        }.get(protocol)
        if not request_cls:
            return None
        req = request_cls(load_balancer_id=lb_id, listener_port=port)
        method = getattr(self.client, f'describe_load_balancer_{protocol}listener_attribute')
        return method(req).body

    def get_health_status(self, lb_id):
        req = slb_models.DescribeHealthStatusRequest(region_id=self.region_id, load_balancer_id=lb_id)
        return self.client.describe_health_status(req).body.backend_servers.backend_server

    def export_to_excel(self, filename='slb_health_info.xlsx'):
        wb = Workbook()
        ws = wb.active
        ws.title = "SLB 健康检查"
        ws.append([
            "负载名称", "实例ID", "负载IP", "监听端口", "协议",
            "监听类型", "监听域名", "URI", "健康检查URI",
            "后端服务器IP", "后端端口"
        ])

        for lb in self.list_load_balancers():
            lb_id = lb.load_balancer_id
            lb_name = lb.load_balancer_name
            listeners, lb_ip = self.get_listener_ports_and_ip(lb_id)
            health_backends = self.get_health_status(lb_id)

            # map: port + instance -> ip + port
            backend_map = {
                f"{b.listener_port}_{b.server_id}": (b.server_ip, b.port)
                for b in health_backends
            }

            for listener in listeners:
                port = listener.listener_port
                proto = listener.listener_protocol.lower()
                attr = self.get_listener_attr(lb_id, port, proto)

                # 默认监听（无转发策略）
                health_uri = getattr(attr, 'health_check_uri', '') or getattr(attr, 'health_check', '') or '未开启'

                for key, (server_ip, server_port) in backend_map.items():
                    if key.startswith(f"{port}_"):
                        # 添加默认监听信息
                        ws.append([
                            lb_name, lb_id, lb_ip, port, proto.upper(),
                            "默认监听", '', '', health_uri,
                            server_ip, server_port
                        ])

                # 查询是否有转发策略
                try:
                    rules_req = slb_models.DescribeRulesRequest(load_balancer_id=lb_id, listener_port=port)
                    rules_resp = self.client.describe_rules(rules_req)
                    rules = rules_resp.body.rules.rule if rules_resp.body.rules else []
                except Exception:
                    rules = []

                # 如果有转发策略，分别列出每条策略
                if rules:
                    for rule in rules:
                        rule_domain = getattr(rule, 'domain', '')
                        rule_uri = getattr(rule, 'url', '')
                        rule_health_uri = getattr(rule, 'health_check_uri', '未开启')
                        vgroup_id = getattr(rule, 'vserver_group_id', '')

                        backend_servers = []
                        if vgroup_id:
                            v_req = slb_models.DescribeVServerGroupAttributeRequest(vserver_group_id=vgroup_id)
                            v_resp = self.client.describe_vserver_group_attribute(v_req)
                            backend_servers = v_resp.body.backend_servers.backend_server if v_resp.body.backend_servers else []

                        if backend_servers:
                            for backend in backend_servers:
                                key = f"{port}_{backend.server_id}"
                                server_ip, server_port = backend_map.get(key, ('', backend.port))
                                ws.append([
                                    lb_name, lb_id, lb_ip, port, proto.upper(),
                                    "转发规则", rule_domain, rule_uri, rule_health_uri,
                                    server_ip, server_port
                                ])

        wb.save(filename)
        print(f"✅ 已导出到 {filename}")

if __name__ == '__main__':
    exporter = SLBExporter()
    exporter.export_to_excel()
